import cv2
import numpy as np
import os
import shutil
import re

def is_patch_valid(patch):
    """判断补丁是否有效，这里假设非全零即为有效"""
    return np.any(patch != 0)

# def get_patch(img, row, col):
#     start_row = row * 64
#     end_row = (row + 1) * 64
#     start_col = col * 64
#     end_col = (col + 1) * 64
#
#     return img[start_row:end_row, start_col:end_col]

def get_patch(img, id):
    row = id // 16
    col = id % 16

    start_row = row * 64
    end_row = (row + 1) * 64
    start_col = col * 64
    end_col = (col + 1) * 64

    return img[start_row:end_row, start_col:end_col]

def increment_filename_number(filename, increment=1):
    """
    将文件名中的数字部分增加 increment，并返回新文件名。
    例如：
        "patches0393.bmp" +1 → "patches0394.bmp"
        "file123.txt" +10 → "file133.txt"
    """
    # 使用正则表达式提取数字部分
    match = re.search(r'(\d+)', filename)
    if not match:
        return filename  # 如果没有数字，直接返回原文件名

    # 提取数字并转换为整数
    number_str = match.group(1)
    number = int(number_str)

    # 数字增加 increment，并保持原位数补零
    new_number = number + increment
    new_number_str = f"{new_number:0{len(number_str)}d}"  # 补零保持位数

    # 替换原数字并生成新文件名
    new_filename = filename.replace(number_str, new_number_str)
    return new_filename


data_dirs = [r'D:\project\FeatureExtration\datasets\trevi', r'D:\project\FeatureExtration\datasets\\notredame']
output_dir = r'D:\project\FeatureExtration\datasets\output'

for id_dir in range(len(data_dirs)):
    info_src_path = os.path.join(data_dirs[id_dir], 'info.txt')
    info_src = np.loadtxt(info_src_path)

    info_dst_path = os.path.join(output_dir, 'info.txt')
    not_first_creat = os.path.exists(info_dst_path)
    if not_first_creat:
        info_dst = np.loadtxt(info_dst_path)
        point_id_max = info_dst[:, 0].max()
        info_src[:, 0] += (point_id_max + 1)

        info_det = np.vstack((info_dst, info_src))
    else:
        shutil.copy(info_src_path, info_dst_path)

    if not_first_creat:
        src_img_files = sorted([f for f in os.listdir(data_dirs[id_dir]) if f.endswith(".bmp")])
        dst_img_files = sorted([f for f in os.listdir(output_dir) if f.endswith(".bmp")])

        src_patch_id = 0
        dst_patch_id = 0
        last_img_file = dst_img_files[-1]
        last_img = cv2.imread(os.path.join(output_dir, last_img_file), cv2.IMREAD_GRAYSCALE)
        for i in range(256):
            if not is_patch_valid(get_patch(last_img, i)):
                dst_patch_id = i
                break

        finished = False
        while not finished:
            dst_img_files = sorted([f for f in os.listdir(output_dir) if f.endswith(".bmp")])
            if dst_patch_id == 0:
                dst_img = np.zeros((1024, 1024), dtype=np.uint8)
                dst_img_file = os.path.join(output_dir, increment_filename_number(dst_img_files[-1]))
            else:
                dst_img = cv2.imread(os.path.join(output_dir, dst_img_files[-1]), cv2.IMREAD_GRAYSCALE)
                dst_img_file = os.path.join(output_dir, dst_img_files[-1])

            src_img = cv2.imread(os.path.join(data_dirs[id_dir], src_img_files[0]), cv2.IMREAD_GRAYSCALE)
            while dst_patch_id <= 255:
            # for i in range(dst_patch_id, 256):
                if src_patch_id >= 256:
                    src_patch_id = 0
                    src_img_files.pop(0)
                    if src_img_files:
                        src_img = cv2.imread(os.path.join(data_dirs[id_dir], src_img_files[0]), cv2.IMREAD_GRAYSCALE)
                    else:
                        finished = True
                        break

                row = dst_patch_id // 16
                col = dst_patch_id % 16
                start_row = row * 64
                end_row = (row + 1) * 64
                start_col = col * 64
                end_col = (col + 1) * 64

                dst_img[start_row:end_row, start_col:end_col] = get_patch(src_img, src_patch_id)
                src_patch_id += 1
                dst_patch_id += 1

            cv2.imwrite(dst_img_file, dst_img)
            print(f"save: {dst_img_file}")
            dst_patch_id = 0
    else:
        src_img_files = sorted([f for f in os.listdir(data_dirs[id_dir]) if f.endswith(".bmp")])
        for img_file in src_img_files:
            shutil.copy(os.path.join(data_dirs[id_dir], img_file), output_dir)
            print(f"copy: {os.path.join(data_dirs[id_dir], img_file)}")

    #
    # if i == 0:
    #     info_file = os.path.join(data_dirs[i], 'info.txt')
    #     info = np.loadtxt(info_file, dtype=np.int64)
    #     img_files = sorted([f for f in os.listdir(data_dirs[i]) if f.endswith(".bmp")])
    #
    #     # 复制文件并重命名
    #     shutil.copy(info_file, os.path.join(output_dir, 'info.txt'))
    #     for img_file in img_files:
    #         shutil.copy(os.path.join(data_dirs[i], img_file), output_dir)
    # else:
    #     info = np.loadtxt(os.path.join(output_dir,'info.txt'), dtype=np.int64)
    #     img_files = sorted([f for f in os.listdir(output_dir) if f.endswith(".bmp")])
    #     max_point_id = np.max(info[:, 0])
    #
    #     src_info_file = os.path.join(data_dirs[i],'info.txt')
    #     src_info = np.loadtxt(src_info_file, dtype=np.int64)
    #     src_img_files = sorted([f for f in os.listdir(data_dirs[i]) if f.endswith(".bmp")])
    #     src_info[:, 0] += (max_point_id+1)
    #     src_info[:, 1] += (len(img_files) * 256)
    #
    #     info = np.vstack((info, src_info))
    #     np.savetxt(os.path.join(output_dir,'info.txt'), info, fmt="%d")
    #
    #     for img_file in src_img_files:
    #         shutil.copy(os.path.join(data_dirs[i], img_file), os.path.join(output_dir, increment_filename_number(img_file, len(img_files))))
    #
